package com.mrx.xgateway.service;

import com.mrx.xgateway.model.Config;
import com.mrx.xgateway.utils.ConfigUtils;
import com.mrx.xgateway.utils.URLUtils;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import kotlin.Pair;
import lombok.SneakyThrows;
import okhttp3.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpMethod;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.util.AntPathMatcher;
import org.springframework.util.PathMatcher;

import java.io.IOException;
import java.util.Collections;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author Mr.X
 * @since 2024-03-10 下午 8:08:53
 */
@Service
public class FrowardService {

    private static final OkHttpClient client = new OkHttpClient.Builder()
            .addInterceptor(FrowardService::logInterceptor)
            .build();
    private static final PathMatcher pathMatcher = new AntPathMatcher();
    private static final Logger logger = LoggerFactory.getLogger(FrowardService.class);
    private static final Set<String> ignoreHeaders = Set.of("host");

    public Object forward(HttpServletRequest request, HttpServletResponse response) {
        Map<String, Config> forwardMap = ConfigUtils.loadConfig().stream()
                .collect(Collectors.toMap(Config::getPath, Function.identity()));
        String requestURI = request.getRequestURI();
        for (Map.Entry<String, Config> entry : forwardMap.entrySet()) {
            String pattern = entry.getKey();
            if (pathMatcher.match(pattern, requestURI)) {
                Config config = entry.getValue();
                String forwardURL = config.getForward();
                if (!config.getStripPath()) {
                    forwardURL = URLUtils.mergeUrl(forwardURL, requestURI);
                } else {
                    String suffixWithoutPattern = pathMatcher.extractPathWithinPattern(pattern, requestURI);
                    forwardURL = URLUtils.mergeUrl(forwardURL, suffixWithoutPattern);
                }
                logger.info("forward: {} -> {}", requestURI, forwardURL);
                return doForward(request, forwardURL, Boolean.TRUE.equals(config.getRequestMethodUppercase()));
            }
        }
        response.setStatus(HttpStatus.NOT_FOUND.value());
        return "No Matched Route Found";
    }

    @SneakyThrows
    private ResponseEntity<?> doForward(HttpServletRequest request, String forward, boolean methodUpperCase) {
        byte[] body = request.getInputStream().readAllBytes();
        Request.Builder requestBuilder = new Request.Builder();
        for (String headerName : Collections.list(request.getHeaderNames())) {
            if (ignoreHeaders.contains(headerName.toLowerCase())) {
                continue;
            }
            for (String headerValue : Collections.list(request.getHeaders(headerName))) {
                requestBuilder.addHeader(headerName, headerValue);
            }
        }
        String method = request.getMethod();
        method = methodUpperCase ? method.toUpperCase() : method.toLowerCase();
        requestBuilder.method(method, HttpMethod.GET.matches(method) ? null : RequestBody.create(body));
        requestBuilder.url(forward);
        try (Response response = client.newCall(requestBuilder.build()).execute()) {
            ResponseEntity.BodyBuilder responseBuilder = ResponseEntity.status(response.code());
            for (Pair<? extends String, ? extends String> header : response.headers()) {
                responseBuilder.header(header.getFirst(), header.getSecond());
            }
            return responseBuilder.body(Optional.ofNullable(response.body()).map(this::parseResponseBody).orElse(null));
        }
    }

    @SneakyThrows
    private byte[] parseResponseBody(ResponseBody responseBody) {
        return responseBody.bytes();
    }

    private static Response logInterceptor(Interceptor.Chain chain) throws IOException {
        Request request = chain.request();
        logger.info("{} {}\n{}", request.method(), request.url(), request.headers());
        return chain.proceed(request);
    }

}
